import argparse
import open3d as o3d
import os
import numpy as np
from tqdm import tqdm
from vis_utils import get_draw_box

lidar2camera = np.array(
    [[-0.007980200000000000, -0.999854100000000000, 0.015104900000000000, 0.151000000000000000],
     [0.118497000000000000, -0.015944500000000000, -0.992826400000000000, -0.461000000000000000],
     [0.992922400000000000, -0.006133100000000000, 0.118606900000000000, -0.915000000000000000],
     [0, 0, 0, 1]])

radar2camera = np.array(
    [[-0.013857, -0.9997468, 0.01772762, 0.05283124],
     [0.10934269, -0.01913807, -0.99381983, 0.98100483],
     [0.99390751, -0.01183297, 0.1095802, 1.44445002],
     [0, 0, 0, 1]])


def get_lidar(lidar_file):
    assert os.path.exists(lidar_file)
    # 修改1  点数变为6
    # return np.fromfile(str(lidar_file), dtype=np.float32).reshape(-1, 6)
    points_lidar = np.fromfile(str(lidar_file), dtype=np.float32).reshape(-1, 4)
    # trans_z = np.array(
    #     [[0.0, 1.0, 0.0, 0.0],
    #      [-1.0, 0.0, 0.0, 0.0],
    #      [0.0, 0.0, 1.0, 0.0],
    #      [0, 0, 0, 1]])
    trans_z = np.array(
        [[1.0, 0.0, 0.0, 0.0],
         [0.0, 1.0, 0.0, 0.0],
         [0.0, 0.0, 1.0, 0.0],
         [0, 0, 0, 1]])
    xyz_lidar = points_lidar[:, :3]
    xyz_lidar1 = np.concatenate((xyz_lidar, np.ones((len(xyz_lidar), 1))), axis=1)
    points_lidar[:, :3] = (trans_z @ xyz_lidar1.T).T[:, :3]
    return points_lidar


def get_ars(ars_file):
    assert os.path.exists(ars_file)
    # 修改1  点数变为5
    points_ars548 = np.fromfile(str(ars_file), dtype=np.float32).reshape(-1, 7)
    num_point = points_ars548.shape[0]
    points_ars548_hom = np.hstack((points_ars548[:, 0:3], np.ones((num_point, 1))))
    # ARS5482V = np.array(
    #     [0.08065852999237734, -0.996741790402992, 6.930878282577699e-05, 0.07152088247656341, 0.9955755395152072,
    #      0.0805607912438313, -0.04836635228309268, 1.9685450015593613, 0.04820318099952712, 0.003970161005102628,
    #      0.9988296607345621, -1.1990909999923147]).reshape([3, 4])
    radar2lidar = np.linalg.inv(lidar2camera) @ radar2camera
    point_radar = np.dot(points_ars548_hom, np.transpose(radar2lidar))
    # points_ars548[:, 0:3] = point_radar
    # trans_z = np.array(
    #     [[0.0, 1.0, 0.0, 0.0],
    #      [-1.0, 0.0, 0.0, 0.0],
    #      [0.0, 0.0, 1.0, 0.0],
    #      [0, 0, 0, 1]])
    trans_z = np.array(
        [[1.0, 0.0, 0.0, 0.0],
         [0.0, 1.0, 0.0, 0.0],
         [0.0, 0.0, 1.0, 0.0],
         [0, 0, 0, 1]])
    # xyz_radar1 = np.concatenate((point_radar, np.ones((len(point_radar), 1))), axis=1)
    xyz_radar1 = point_radar
    points_ars548[:, 0:3] = (trans_z @ xyz_radar1.T).T[:, :3]

    return points_ars548


def get_arbe(arbe_file):
    assert os.path.exists(arbe_file)
    # 修改1  点数变为5
    points_arbe = np.fromfile(str(arbe_file), dtype=np.float32).reshape(-1, 5)
    num_point = points_arbe.shape[0]
    points_arbe_hom = np.hstack((points_arbe[:, 0:3], np.ones((num_point, 1))))
    ARB2V = np.array([0.9981128011677526, 0.05916115557244023, 0.016455814060541557, 0.07800451346438697,
                      -0.059503891609836816, 0.9980033119043885, -0.021181980812864695, 2.214080041485726,
                      0.015169806470300943, 0.022121191179620064, 0.9996402002082792, -1.6030740415943632]).reshape(
        [3, 4])
    point_radar = np.dot(points_arbe_hom, np.transpose(ARB2V))
    # points_arbe[:,0:3] = point_lidar
    trans_z = np.array(
        [[0.0, 1.0, 0.0, 0.0],
         [-1.0, 0.0, 0.0, 0.0],
         [0.0, 0.0, 1.0, 0.0],
         [0, 0, 0, 1]])
    xyz_radar1 = np.concatenate((point_radar, np.ones((len(point_radar), 1))), axis=1)
    points_arbe[:, 0:3] = (trans_z @ xyz_radar1.T).T[:, :3]

    return points_arbe


def get_boxes(label_file):
    assert os.path.exists(label_file)
    objs = []
    with open(label_file, 'r') as f:
        # open为打开文件，r为读取
        f = open(label_file, 'r')
        # 逐行读取文件内容
        lines = f.readlines()
        for line in lines:
            # 消除空格和空行
            line = line.strip()
            tokens = line.split(' ')
            objs.append(tokens)

    boxes = np.empty(shape=(0, 7))

    # 预测框坐标系变换：相机坐标系--->lidar坐标系
    # tr_velo2cam = np.array([[0.9994015951319486, -0.03199927972993585, 0.013133839670170688, 0.2710357611083472],
    #                         [0.013650110615193858, 0.015969147292999922, -0.9997793060545548, -0.11226274034223951],
    #                         [0.03178248146242525, 0.9993603116151348, 0.01639638498544132, -0.02040423686369943],
    #                         [0, 0, 0, 1]])
    # tr_cam2velo = tr_velo2cam.copy()
    # tr_cam2velo[:3, :3] = np.linalg.inv(tr_velo2cam[:3, :3])
    # tr_cam2velo[:3, -1] = np.dot(-tr_cam2velo[:3, :3], tr_velo2cam[:3, -1])

    tr_cam2velo = np.linalg.inv(lidar2camera)

    for obj in objs:
        # box = np.array([[float(obj[11]), float(obj[12]), float(obj[13]), float(obj[10]), float(obj[9]), float(obj[8]),
        #                  -float(obj[14])]])
        box = np.array([[float(obj[11]), float(obj[12]), float(obj[13]), float(obj[9]), float(obj[10]), float(obj[8]),
                         -float(obj[14])]])
        # 3D框中心点转换坐标--->
        # trans_z = np.array(
        #     [[0.0, 1.0, 0.0, 0.0],
        #      [-1.0, 0.0, 0.0, 0.0],
        #      [0.0, 0.0, 1.0, 0.0],
        #      [0, 0, 0, 1]])
        trans_z = np.array(
            [[1.0, 0.0, 0.0, 0.0],
             [0.0, 1.0, 0.0, 0.0],
             [0.0, 0.0, 1.0, 0.0],
             [0, 0, 0, 1]])
        box_center = box[:, :3]
        box_center1 = np.concatenate((box_center, np.ones((len(box_center), 1))), axis=1)
        box[:, :3] = (trans_z @ tr_cam2velo @ box_center1.T).T[:, :3]
        box[:, 2] += 0.5 * float(obj[8])  # 高度修正，kitti相机坐标系以地面为中心 lidar坐标系通常以中心点为中心

        boxes = np.concatenate((boxes, box))
    return boxes


def main(args):
    assert os.path.exists(args.data_path)
    datasets = os.listdir(args.data_path)
    datasets.sort(key=lambda x: x)
    for dataset in tqdm(datasets):
        if "__" == dataset[:2]:  # 数据集以"__"开头
            dataset_path = os.path.join(args.data_path, dataset)
            main_worker(args, dataset_path)


def main_worker(args, dataset_path):
    lidar_size = len(os.listdir(dataset_path + '/lidar_bin'))
    radar_size = len(os.listdir(dataset_path + '/radar_bin'))
    label_size = len(os.listdir(dataset_path + '/samples/label_txt'))
    assert (lidar_size == radar_size and lidar_size >= label_size)
    names = []
    for name in os.listdir(dataset_path + '/samples/label_txt'):
        names.append(str.split(name, '.')[0])

    for name in names:

        lidar_path = dataset_path + '/lidar_bin/' + name + ".bin"
        radar_path_ars548 = dataset_path + '/radar_bin/' + name + ".bin"
        # radar_path_arbe = dataset_path + '/radar_bin_arbe/' + name + ".bin"
        label_path = dataset_path + '/samples/label_txt/' + name + ".txt"

        lidar = get_lidar(lidar_path)
        radar_ars548 = get_ars(radar_path_ars548)
        # radar_arbe = get_arbe(radar_path_arbe)

        vis = o3d.visualization.Visualizer()
        vis.create_window(window_name="show_pred_pcd")
        render_option = vis.get_render_option()
        render_option.point_size = 2
        coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=2.0, origin=[0, 0, 0])
        vis.add_geometry(coordinate_frame)

        datas = o3d.geometry.PointCloud()
        datas.points = o3d.utility.Vector3dVector(lidar[:, :3])
        datas.paint_uniform_color([0.6, 0.6, 0.6])
        vis.add_geometry(datas)

        data_radar_ars548 = o3d.geometry.PointCloud()
        data_radar_ars548.points = o3d.utility.Vector3dVector(radar_ars548[:, :3])
        data_radar_ars548.paint_uniform_color([1, 0, 0])
        vis.add_geometry(data_radar_ars548)

        # data_radar_arbe = o3d.geometry.PointCloud()
        # data_radar_arbe.points = o3d.utility.Vector3dVector(radar_arbe[:, :3])
        # data_radar_arbe.paint_uniform_color([0, 0, 1])
        # vis.add_geometry(data_radar_arbe)

        boxes = get_boxes(label_path)
        exp_draw_boxes = get_draw_box(boxes)
        for box in exp_draw_boxes:
            vis.add_geometry(box)

        vis.run()
        vis.destroy_window()

    print(0)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/files/data/view_of_delft_PUBLIC",
                        help='your data root for kitti')

    args = parser.parse_args()

    main(args)
